from tensorflow.python.framework.graph_util import convert_variables_to_constants
from keras.layers.normalization import BatchNormalization
from keras.layers.advanced_activations import LeakyReLU
from keras.regularizers import l2
from keras.models import Model,Input
import json
from keras.layers import Conv2D, MaxPool2D, Add, ZeroPadding2D, UpSampling2D, Concatenate
import tensorflow as tf
from tensorflow.python.framework import graph_io
from keras import backend as K
import os


def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):

    graph = session.graph
    with graph.as_default():
        freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
        output_names = output_names or []
        output_names += [v.op.name for v in tf.global_variables()]
        print("output_names",output_names)
        input_graph_def = graph.as_graph_def()
        print("input_graph_def:",input_graph_def)
        if clear_devices:
            for node in input_graph_def.node:
                node.device = ""
        frozen_graph = convert_variables_to_constants(session, input_graph_def, output_names, freeze_var_names)
    return frozen_graph

def darknet(input):
    network = NetworkConv2D_BN_Leaky(input=input, channels=32, kernel_size=(3,3))
    network = residual_block(input=network, channels=64, num_blocks=1)
    network = residual_block(input=network, channels=128, num_blocks=2)
    network = residual_block(input=network, channels=256, num_blocks=8)
    network = residual_block(input=network, channels=512, num_blocks=8)
    network = residual_block(input=network, channels=1024, num_blocks=4)


    return network

def residual_block(input, channels, num_blocks):
    network = ZeroPadding2D(((1,0), (1,0)))(input)
    network = NetworkConv2D_BN_Leaky(input=network,channels=channels, kernel_size=(3,3), strides=(2,2), padding="valid")

    for blocks in range(num_blocks):
        network_1 = NetworkConv2D_BN_Leaky(input=network, channels= channels // 2, kernel_size=(1,1))
        network_1 = NetworkConv2D_BN_Leaky(input=network_1,channels= channels, kernel_size=(3,3))

        network = Add()([network, network_1])
    return network

def last_layers(input, channels_in, channels_out, layer_name=""):



    network = NetworkConv2D_BN_Leaky( input=input, channels=channels_in, kernel_size=(1,1))
    network = NetworkConv2D_BN_Leaky(input=network, channels= (channels_in * 2) , kernel_size=(3, 3))
    network = NetworkConv2D_BN_Leaky(input=network, channels=channels_in, kernel_size=(1, 1))
    network = NetworkConv2D_BN_Leaky(input=network, channels=(channels_in * 2), kernel_size=(3, 3))
    network = NetworkConv2D_BN_Leaky(input=network, channels=channels_in, kernel_size=(1, 1))

    network_1 = NetworkConv2D_BN_Leaky(input=network, channels=(channels_in * 2), kernel_size=(3, 3))
    network_1 = Conv2D(filters=channels_out, kernel_size=(1,1), name=layer_name)(network_1)

    return  network, network_1


def NetworkConv2D_BN_Leaky(input, channels, kernel_size, kernel_regularizer = l2(5e-4), strides=(1,1), padding="same", use_bias=False):

    network = Conv2D( filters=channels, kernel_size=kernel_size, strides=strides, padding=padding, kernel_regularizer=kernel_regularizer, use_bias=use_bias)(input)
    network = BatchNormalization()(network)
    network = LeakyReLU(alpha=0.1)(network)
    return network




def h_to_pb(configuration_json,detection_model_path):
    class LoadModel:
        def __init__(self, configuration_json, detection_model_path):
            self.__model_type = "yolov3"
            self.__detection_config_json_path = configuration_json
            self.__model_path = detection_model_path

        def loadModel(self):
            """
            'loadModel' is used to load the model into the CustomObjectDetection class
            :return: None
            """

            def yolo_main(input, num_anchors, num_classes):
                darknet_network = Model(input, darknet(input))
                network, network_1 = last_layers(darknet_network.output, 512, num_anchors * (num_classes + 5),
                                                 layer_name="last1")
                network = NetworkConv2D_BN_Leaky(input=network, channels=256, kernel_size=(1, 1))
                network = UpSampling2D(2)(network)
                network = Concatenate()([network, darknet_network.layers[152].output])
                network, network_2 = last_layers(network, 256, num_anchors * (num_classes + 5), layer_name="last2")
                network = NetworkConv2D_BN_Leaky(input=network, channels=128, kernel_size=(1, 1))
                network = UpSampling2D(2)(network)
                network = Concatenate()([network, darknet_network.layers[92].output])
                network, network_3 = last_layers(network, 128, num_anchors * (num_classes + 5), layer_name="last3")
                return Model(input, [network_1, network_2, network_3])

            if self.__model_type == "yolov3":
                detection_model_json = json.load(open(self.__detection_config_json_path,'rb'))
                self.__model_labels = detection_model_json["labels"]
                self.__model_anchors = detection_model_json["anchors"]
                self.__model = yolo_main(Input(shape=(None, None, 3), name='input_img'), 3, len(self.__model_labels))
                self.__model.load_weights(self.__model_path)
                return self.__model

    graph = tf.Graph()  # 声明tf.Graph()的一个类实例，即获取一个graph
    with graph.as_default():
        LoadModel = LoadModel(configuration_json, detection_model_path)
        model = LoadModel.loadModel()
        frozen_graph = freeze_session(K.get_session(), output_names=[out.op.name for out in model.outputs])
        print([out.op.name for out in model.outputs])
        path,model_name = os.path.split(detection_model_path)
        pb_model_name = model_name.replace(".h5",".pb")
        graph_io.write_graph(frozen_graph, path, pb_model_name, as_text=False)
        return os.path.join(path, pb_model_name)